# -*- coding: utf-8 -*-
"""
轮动策略框架 | 邢不行 | 2024分享会
author: 邢不行
微信: xbx6660
"""
from Config import *
import pandas as pd
from plotly.offline import plot
from plotly.subplots import make_subplots
import plotly.graph_objs as go
import plotly.express as px
from Functions import revise_data_length


def draw_equity_curve_plotly1(df, data_dict, date_col=None, right_axis=None, pic_size=[1500, 800], chg=False,
                             title=None, path=root_path + '/data/pic1.html', show=True):
    """
    绘制策略曲线
    :param df: 包含净值数据的df
    :param data_dict: 要展示的数据字典格式：｛图片上显示的名字:df中的列名｝
    :param date_col: 时间列的名字，如果为None将用索引作为时间列
    :param right_axis: 右轴数据 ｛图片上显示的名字:df中的列名｝
    :param pic_size: 图片的尺寸
    :param chg: datadict中的数据是否为涨跌幅，True表示涨跌幅，False表示净值
    :param title: 标题
    :param path: 图片路径
    :param show: 是否打开图片
    :return:
    """
    draw_df = df.copy()

    # 设置时间序列
    if date_col:
        time_data = draw_df[date_col]
    else:
        time_data = draw_df.index

    # 绘制左轴数据
    fig = make_subplots(specs=[[{"secondary_y": True}]])
    for key in data_dict:
        if chg:
            draw_df[data_dict[key]] = (draw_df[data_dict[key]] + 1).fillna(1).cumprod()
        fig.add_trace(go.Scatter(x=time_data, y=draw_df[data_dict[key]], name=key))

    # 绘制右轴数据
    if right_axis:
        key = list(right_axis.keys())[0]
        fig.add_trace(go.Scatter(x=time_data, y=draw_df[right_axis[key]], name=key + '(右轴)',
                                 #  marker=dict(color='rgba(220, 220, 220, 0.8)'),
                                 opacity=0.1, line=dict(width=0),
                                 marker_color='orange',
                                 fill='tozeroy',
                                 yaxis='y2'))  # 标明设置一个不同于trace1的一个坐标轴


    fig.update_layout(template="none", width=pic_size[0], height=pic_size[1], title_text=title,
                      hovermode="x unified", hoverlabel=dict(bgcolor='rgba(255,255,255,0.5)', ),
                      )
    fig.update_layout(
        updatemenus=[
            dict(
                buttons=[
                    dict(label="线性 y轴",
                         method="relayout",
                         args=[{"yaxis.type": "linear"}]),
                    dict(label="Log y轴",
                         method="relayout",
                         args=[{"yaxis.type": "log"}]),
                ])],
    )
    plot(figure_or_data=fig, filename=path, auto_open=False)

    fig.update_yaxes(
        showspikes=True, spikemode='across', spikesnap='cursor', spikedash='solid', spikethickness=1,  # 峰线
    )
    fig.update_xaxes(
        showspikes=True, spikemode='across+marker', spikesnap='cursor', spikedash='solid', spikethickness=1,  # 峰线
    )

    # 打开图片的html文件，需要判断系统的类型
    if show:
        res = os.system('start %s' % path)
        if res != 0:
            os.system('open %s' % path)


def plot_ratation_net_value1(equity, factor_list, rtn):
    '''
        画轮动资金曲线图1
    :param equity: 计算得到的资金曲线数据
    :param factor_list: 因子列表
    :param rtn: 策略评价指标
    :return:
    '''
    # 合并BTC资金曲线
    BTC = pd.read_csv(swap_path + 'BTC-USDT.csv', encoding='gbk', parse_dates=['candle_begin_time'], skiprows=1)
    BTC['BTC涨跌幅'] = BTC['close'].pct_change()
    equity = pd.merge(left=equity, right=BTC[['candle_begin_time', 'BTC涨跌幅']], on=['candle_begin_time'], how='left')
    equity['BTC涨跌幅'].fillna(value=0, inplace=True)
    equity['BTC资金曲线'] = (equity['BTC涨跌幅'] + 1).cumprod()

    # 指定左轴的数据
    data_dict = {'轮动资金曲线': '轮动资金曲线', 'BTC资金曲线': 'BTC资金曲线'}

    # 指定右轴使用数据
    right_axis = {'轮动回撤曲线': 'dd2here'}

    # 指定画图标题
    pic_title = '轮动策略：factors:%s_nv:%s_pro:%s_risk:%s' % (factor_list, rtn.at['累积净值', 0], rtn.at['年化收益', 0], rtn.at['最大回撤', 0])

    # 调用画图函数
    draw_equity_curve_plotly1(equity, data_dict=data_dict, date_col='candle_begin_time', right_axis=right_axis,
                              title=pic_title)


def draw_equity_curve_plotly2(df, data_dict, date_col=None, right_axis=None, pic_size=[1500, 800], chg=False,
                             title=None, path=root_path + '/data/pic2.html', show=True, color_dict={}):
    """
    绘制策略曲线
    :param df: 包含净值数据的df
    :param data_dict: 要展示的数据字典格式：｛图片上显示的名字:df中的列名｝
    :param date_col: 时间列的名字，如果为None将用索引作为时间列
    :param right_axis: 右轴数据 ｛图片上显示的名字:df中的列名｝
    :param pic_size: 图片的尺寸
    :param chg: datadict中的数据是否为涨跌幅，True表示涨跌幅，False表示净值
    :param title: 标题
    :param path: 图片路径
    :param show: 是否打开图片
    :param color_dict: 各策略对应的颜色
    :return:
    """
    draw_df = df.copy()

    # 设置时间序列
    if date_col:
        time_data = draw_df[date_col]
    else:
        time_data = draw_df.index

    # 绘制左轴数据
    fig = make_subplots(specs=[[{"secondary_y": True}]])
    for key in data_dict:
        if chg:
            draw_df[data_dict[key]] = (draw_df[data_dict[key]] + 1).fillna(1).cumprod()
        fig.add_trace(go.Scatter(x=time_data, y=draw_df[data_dict[key]], name=key, marker_color=color_dict[key]))

    # 绘制右轴数据
    if right_axis:
        key = list(right_axis.keys())[0]
        for label in draw_df['label'].unique():
            begin_index = draw_df[draw_df['label'] == label].index[0]
            end_index = draw_df[draw_df['label'] == label].index[-1]
            temp = draw_df.loc[begin_index - 1: end_index, :]
            fig.add_trace(go.Scatter(x=temp[date_col], y=temp[right_axis[key]], name='轮动资金曲线', line=dict(color=color_dict[temp['symbol'].iloc[-1].strip()]), yaxis='y2'))

    fig.update_layout(template="none", width=pic_size[0], height=pic_size[1], title_text=title,
                      hovermode="x unified", hoverlabel=dict(bgcolor='rgba(255,255,255,0.5)', ),
                      )
    fig.update_layout(
        updatemenus=[
            dict(
                buttons=[
                    dict(label="线性 y轴",
                         method="relayout",
                         args=[{"yaxis.type": "linear"}]),
                    dict(label="Log y轴",
                         method="relayout",
                         args=[{"yaxis.type": "log"}]),
                ])],
    )
    plot(figure_or_data=fig, filename=path, auto_open=False)

    fig.update_yaxes(
        showspikes=True, spikemode='across', spikesnap='cursor', spikedash='solid', spikethickness=1,  # 峰线
    )
    fig.update_xaxes(
        showspikes=True, spikemode='across+marker', spikesnap='cursor', spikedash='solid', spikethickness=1,  # 峰线
    )

    # 打开图片的html文件，需要判断系统的类型
    if show:
        res = os.system('start %s' % path)
        if res != 0:
            os.system('open %s' % path)


def plot_ratation_net_value2(equity, base_equity, factor_list, rtn):
    '''
        画轮动资金曲线2
    :param equity: 资金曲线数据
    :param base_equity: 读取到的子策略资金曲线数据
    :param factor_list: 因子列表
    :param rtn: 策略评价指标
    :return:
    '''
    # 整理子策略资金曲线数据
    base_equity = base_equity[['candle_begin_time', 'symbol', '每小时涨跌幅']]
    base_equity = base_equity[base_equity['candle_begin_time'] >= equity['candle_begin_time'].iloc[0]].reset_index(drop=True)
    _equity = equity[['candle_begin_time']]
    for symbol in base_equity['symbol'].unique():
        _equity[symbol] = revise_data_length(base_equity[base_equity['symbol'] == symbol]['每小时涨跌幅'].sum(), len(equity))
        _equity[symbol] = (_equity[symbol] + 1).cumprod()

    # =将轮动策略的资金曲线补充到之前的_equity中，以及添加label
    _equity['轮动资金曲线'] = equity['轮动资金曲线'].values.tolist()
    _equity['symbol'] = equity['策略名称'].values.tolist()
    _equity['label'] = equity['label'].values.tolist()

    # 指定左轴的数据
    data_dict = {}
    for col in base_equity['symbol'].unique():
        data_dict[col] = col

    # 指定右轴使用数据
    right_axis = {'轮动资金曲线': '轮动资金曲线'}

    # 指定左轴使用的颜色集
    color_list = ['red', 'green', 'blue', 'orange', 'purple', 'pink', 'teal', 'olive', 'navy', 'maroon', 'cyan',
                  'magenta', 'violet', 'silver', 'gray', 'tan', 'coral', 'turquoise', 'indigo', 'salmon']

    # 针对各个左轴的资金曲线配对相应颜色
    color_dict = {'轮动资金曲线': '轮动资金曲线'}
    for i, col in enumerate(data_dict):
        color_dict[col] = color_list[i]

    # =指定画图标题
    pic_title = '轮动策略：factors:%s_nv:%s_pro:%s_risk:%s' % (factor_list, rtn.at['累积净值', 0], rtn.at['年化收益', 0], rtn.at['最大回撤', 0])

    # =调用画图函数
    draw_equity_curve_plotly2(_equity, data_dict=data_dict, date_col='candle_begin_time', right_axis=right_axis,
                              title=pic_title, color_dict=color_dict)


def draw_bar_plotly(data, x, y, title='分箱图', text='asset', pic_size=[1200, 500], show=True, path=root_path + '/data/bar_pic.html'):
    # 配置fig的列表，生成图之后先添加到fig_list中，最后一起画图
    bar_fig = px.bar(data, x=x, y=y, title=title, text=text, width=pic_size[0], height=pic_size[1])
    plot(figure_or_data=bar_fig, filename=path, auto_open=False)
    # 打开图片的html文件，需要判断系统的类型
    if show:
        res = os.system('start %s' % path)
        if res != 0:
            os.system('open %s' % path)
